import os

import cv2 as cv
import torch
from torch.utils.data import Dataset

class CodeDataset(Dataset):

    def __init__(self, path='', transform=None):#, transform=None
        self.path = path
        self.transform = transform
        self.imgs = []
        for root, dirs, files in os.walk(path):
            for file in files:
                self.imgs.append(file)

        self.n_samples = self.imgs.__len__()
        self.labels = [i.split('.')[0] for i in self.imgs]
        unique_set=set(char for label in self.labels for char in label)
        self.map=sorted(list(unique_set))
        self.characters = ['-']+sorted(list(unique_set))
        self.myclass_len=len(self.characters)
        self.char_to_num = dict((c, i) for i, c in enumerate(self.characters))
        self.num_to_char = dict((i, c) for i, c in enumerate(self.characters))


    def __getitem__(self, index):
        img_path = os.path.join(self.path, self.imgs[index])
        to_img = cv.imread(img_path)
        if self.transform:
            to_img = self.transform(to_img)
        label = [self.char_to_num[i] for i in self.labels[index]]
        return to_img, torch.tensor(label,dtype=torch.int32)

    def __len__(self):
        return self.n_samples

if __name__ == '__main__':
    dataset = CodeDataset(path=r'C:\Users\Administrator\Desktop\WangTing\five_numbers_verification_code')
    train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=4, shuffle=True, num_workers=0)


